from numpy import *
from scipy import optimize
import ConfigParser
import scientific_utils as su
try:
	import gdspy
	reload(gdspy)
except ImportError:
	print 'gdspy not found, gds exporting disabled.'
	gdspy = 0
try:
	import matplotlib.pyplot as plt
except ImportError:
	print 'matplotlib not found, visualization methods disabled'
	plt = 0



class Spectrometer:
	def __init__(self):
		self.changed = True

#	def setInputWithThisWavelengthReflectingBack(self,keff):
#		self.changed = True
#		angle = arcsin(self.order*2*pi/keff/self.grating.pitch/2)
#		self.input.angle = angle + pi
#		self.input.point = array([cos(2*angle),sin(2*angle)])*self.radius
		
#	def makeLineAlongWidth(self,angle,width=None,center=None,N_points=None):
#		aperture = Aperture()
#		if center == None:
#			center = self.radius*array([cos(2*angle),sin(2*angle)])
#		aperture.center = center
#		if width == None:
#			width = self.output.width
#		aperture.width = width
#		if N_points == None:
#			N_points = self.output.n
#		aperture.angle = angle
#		aperture.eval_r2 = True
#		aperture.x = linspace(-width/2,width/2,N_points)
#		aperture.angle = angle + pi/2
#		aperture.points = transpose(vstack((cos(aperture.angle)*aperzture.x,sin(aperture.angle)*aperture.x)))+center
#		aperture.n = N_points
#		return aperture	
		
	def changeApertureAngle(self,aperture,angle):
		aperture.center = self.radius*array([cos(2*angle),sin(2*angle)])
		aperture.angle = angle+pi
		aperture.eval_r2 = True
		aperture.points = transpose(vstack((cos(angle+pi/2)*aperture.x,sin(angle+pi/2)*aperture.x)))+aperture.center
		return
		
	def setCenter(self,aperture):
		angle = (aperture.angle-pi)*2
		aperture.center = self.radius*array([cos(angle),sin(angle)])		
		
	def setRowlandCircleEdgePoints(self,startAngle,endAngle,stepSize,width):
		output = Aperture()
		output.eval_r2 = True
		output.angles = linspace(startAngle,endAngle,self.radius*(endAngle-startAngle)/stepSize)
		output.n = len(output.angles)
		output.points = array([cos(output.angles),sin(output.angles)]).transpose()*self.radius
		output.x = output.angles*180/pi
		output.widths = ones(output.n)*width
		output.angles = output.angles/2+pi
		return output

	def makeGrating(self,aberrationFreeKeff,centralGroove=None):
		self.changed = True
		self.grating.aberrationFreeKeff = aberrationFreeKeff
		self.grating.aberrationFreeAngle = arcsin(sin(self.input.angle)+self.order/self.grating.pitch*2*pi/aberrationFreeKeff)
		self.grating.aberrationFreePoint = array([cos(2*self.grating.aberrationFreeAngle),sin(2*self.grating.aberrationFreeAngle)])*self.radius
		self.grating.aberrationFreeAngle += pi
		if centralGroove == None:
			centralGroove = self.grating.n/2	
		def normalizedPathLength(theta):
			I = self.input.center
			S = self.grating.aberrationFreePoint
			Gx = (1-2*cos(theta))*self.radius 
			Gy = 2*sin(theta)*self.radius
			IG = sqrt((Gx-I[0])**2+(Gy-I[1])**2)	
			GS = sqrt((Gx-S[0])**2+(Gy-S[1])**2)
			return aberrationFreeKeff*(IG+GS)/2/pi
		Ref = normalizedPathLength(0)
		self.grating.points = zeros((self.grating.n,2))
		for i in arange(self.grating.n):
			def g(theta):
				s = normalizedPathLength(theta)
				s = s-self.order*(i-centralGroove)
				s = s-Ref
				return s
			angle = optimize.fsolve(g,0)
			self.grating.points[i,:] = array([(-2*cos(angle)+1),2*sin(angle)])*self.radius	

			
	def makeOneStigmaticPointGrating(self,aberrationFreeKeff,centralGroove=None):
		self.changed = True
		self.grating.aberrationFreeKeff = aberrationFreeKeff
		self.grating.aberrationFreeAngle = arcsin(sin(self.input.angle)+self.order/self.grating.pitch*2*pi/aberrationFreeKeff)
		self.grating.aberrationFreePoint = array([cos(2*self.grating.aberrationFreeAngle),sin(2*self.grating.aberrationFreeAngle)])*self.radius
		self.grating.aberrationFreeAngle += pi
		if centralGroove == None:
			centralGroove = self.grating.n/2	
		def normalizedPathLength(theta):
			I = self.input.center
			S = self.grating.aberrationFreePoint
			Gx = (1-2*cos(theta))*self.radius 
			Gy = 2*sin(theta)*self.radius
			IG = sqrt((Gx-I[0])**2+(Gy-I[1])**2)	
			GS = sqrt((Gx-S[0])**2+(Gy-S[1])**2)
			return aberrationFreeKeff*(IG+GS)/2/pi
		Ref = normalizedPathLength(0)
		self.grating.points = zeros((self.grating.n,2))
		for i in arange(self.grating.n):
			def g(theta):
				s = normalizedPathLength(theta)
				s = s-self.order*(i-centralGroove)
				s = s-Ref
				return s
			angle = optimize.fsolve(g,0)
			self.grating.points[i,:] = array([(-2*cos(angle)+1),2*sin(angle)])*self.radius	

	def blazeGratingToPoint(self,blazePoint):
		self.changed = True
		I = self.input.center-self.grating.points
		mI = sqrt(sum(I**2,1))
		S = blazePoint-self.grating.points
		angleI = su.angle(I)
		self.grating.angles = (angleI+su.angle(S))/2
		an = diff((self.grating.angles[1:]+self.grating.angles[:-1])/2)
		an = hstack((an,an[-1]))
		an = hstack((an[0],an))
		self.grating.widths = an*mI/cos(self.grating.angles-angleI)
		#self.grating.widths = self.grating.width*ones(self.grating.n)

	def setOuputsCollect(wavelengths,transmissionFunction='default'):
		self.changed = True
		output=1		
	
	def plot(self):
		if plt==0:
			print 'matplotlib not loaded'
			return -1
		self.grating.plot()
		self.input.plot()
		if hasattr(self,'output'):
			self.output.plot()
		if hasattr(self,'outputs'):
			for output in self.outputs:
				output.plot()
		t=linspace(0,2*pi,100)
		plt.plot(self.radius*cos(t),self.radius*sin(t))
		plt.axes().set_aspect('equal')
		
	def write(self,filePointer):
		cfg = ConfigParser.ConfigParser()
		cfg.add_section('Parameters')
		cfg.set('Parameters','radius',self.radius)
		cfg.set('Parameters','order',self.order)
		cfg.set('Parameters','aberrationFreeKeff',self.grating.aberrationFreeKeff)
		cfg.add_section('Detail')
		cfg.set('Detail','GratingPoints',repr(self.grating.points)	)
		cfg.set('Detail','GratingAngles',repr(self.grating.angles.view()	))
		cfg.set('Detail','GratingWidths',repr(self.grating.widths.view()	))
		cfg.set('Detail','InputPoint'	,repr(self.input.center.view()		))
		cfg.set('Detail','InputAngle'	,repr(self.input.angle				))
		cfg.set('Detail','InputWidth'	,repr(self.input.width				))
		points = array([output.center for output in self.outputs])
		angles = array([output.angle for output in self.outputs])
		widths = array([output.width for output in self.outputs])
		cfg.set('Detail','OutputPoints' ,repr(points.view()		))
		cfg.set('Detail','OutputAngles'	,repr(angles.view()		))
		cfg.set('Detail','OutputWidths'	,repr(widths.view()		))
		cfg.write(filePointer)
		
# 	def write(self,filePointer):
# 		if gdspy==0:
# 			print 'gdspy not loaded.'
# 			return -1
		
class DiffractionGrating:
	def __init__(self):
		pass
	
	def plot(self):
		if plt==0:
			print 'matplotlib not loaded.'
			return -1
		plt.plot(self.points[:,0],self.points[:,1],'o')
		
	def plotField(self):
		if plt==0:
			print 'matplotlib not loaded.'
			return -1
		plt.plot(abs(self.E))		
	
# 	def gds(self):
# 		if gdspy==0:
# 			print 'gdspy not loaded.'
# 			return -1
# 		cell = gdspy.Cell('DiffractionGrating')
# 		
# 		for i in range(self.nGrooves-1):
# 			points = zeros((5,2))
# 			points[0] = self.points[i] + 0.5*self.widths[i]*array([cos(self.angles[i]  +pi/2),sin(self.angles[i]  +pi/2)])
# 			points[1] = self.points[i] - 0.5*self.widths[i]*array([cos(self.angles[i]  +pi/2),sin(self.angles[i]  +pi/2)])
# 			points[2] = self.points[i+1]+0.5*self.widths[i]*array([cos(self.angles[i+1]+pi/2),sin(self.angles[i+1]+pi/2)])
# 			points[3] = points[2]-array([10e-6,0])
# 			points[4] = points[0]-array([10e-6,0])
# 			cell.add(gdspy.Polygon(2, points))
# 		points = zeros((4,2))
# 		points = self.points[i]+0.5*self.widths[i]*array([cos(self.angles[i]+pi/2),sin(self.angles[i]+pi/2)])
# 		points = self.points[i]-0.5*self.widths[i]*array([cos(self.angles[i]+pi/2),sin(self.angles[i]+pi/2)])
# 		points = points[1]-array([10e-6,0])
# 		gdspy.gds_view()



#class Input:
#	def __init__(self,width):
#		self.changed = True
#		self.width = width
#		self.point=0
#		self.angle=0
#
#	def plot(self):
#		if plt==0:
#			print 'matplotlib no loaded.'
#			return -1
#		plt.plot([self.point[0]],[self.point[1]],'o')
	
class Aperture:
	def __init__(self):
		self.eval_r2 = True
		self.points=0
		self.widths=0
		self.angles=0
		
	def plot(self):
		if plt==0:
			print 'matplotlib not loaded.'
			return -1
		
		if size(self.points)==0:
			return
		plt.plot(self.points[...,0],self.points[...,1],'o')
		vy = array([cos(self.angle),sin(self.angle)])
		vx = array([cos(self.angle-pi/2),sin(self.angle-pi/2)])
		arrow = zeros((8,2))
		arrow[0] = self.center+vx*0.25*self.width
		arrow[1] = self.center-vx*0.25*self.width
		arrow[2] = self.center-vx*0.25*self.width+vy*0.5*self.width
		arrow[3] = self.center-vx*0.50*self.width+vy*0.5*self.width
		arrow[4] = self.center+vy*1*self.width
		arrow[5] = self.center+vx*0.50*self.width+vy*0.5*self.width
		arrow[6] = self.center+vx*0.25*self.width+vy*0.5*self.width
		arrow[7] = arrow[0]
		plt.plot(arrow[...,0],arrow[...,1])

	def plotField(self):
		if plt==0:
			print 'matplotlib not loaded.'
			return -1
		
		if size(self.points)==0:
			return
		plt.plot(self.x,abs(self.E))

class Grid:
	def __init__(self,Xlength,Ylength,pitch,center):
		X = arange(-Xlength/2,Xlength/2,pitch)
		Y = arange(-Ylength/2,Ylength/2,pitch)
		self.eval_r2 = True
		self.nX = len(X)
		self.nY = len(Y)
		self.n = self.nY*self.nX
		m = array(meshgrid(X,Y))
		self.points =  m.flatten().reshape(2,self.n).transpose()+center
	
	def plotField(self):
		if plt==0:
			print 'matplotlib not loaded.'
			return -1

		plt.imshow(flipud(abs(self.E).reshape(self.nY,self.nX)))

	def plot(self):
		if plt==0:
			print 'matplotlib not loaded.'
			return -1
		
		if size(self.points)==0:
			return
		plt.plot(self.points[...,0],self.points[...,1],'o')
		


# 	def gds(self):
# 		cell = gdspy.Cell('CollectionWaveguides')
# 		TaperLength = 50e-6
# 		WaveguideWidth = 500e-9
# 		vec = array([cos(self.angles[0]+pi/2),sin(self.angles[0]+pi/2)])
# 		points = zeros((4,2))
# 		points[0] = self.points[0]+0.5*self.widths[0]*vec
# 		points[1] = points[0]+1e-6*vec
# 		points[2] = points[1]+TaperLength*array([cos(self.angles[0]+pi),sin(self.angles[0]+pi)])
# 		points[3] = points[2]-vec*(1e-6+(self.widths[0]-WaveguideWidth)/2)
# 		RightCollectionPoint = self.points[0]-0.5*self.widths[0]*vec
# 		RightOutputPoint = points[2]-vec*(1e-6+self.widths[0]/2+WaveguideWidth)
# 		cell.add(gdspy.Polygon(2, points))
# 		for i in range(1,self.n-1):
# 			vec = array([cos(self.angles[i]+pi/2),sin(self.angles[i]+pi/2)])
# 			points = zeros((4,2))
# 			points[0] = self.points[i]+0.5*self.widths[i]*vec
# 			points[1] = RightCollectionPoint
# 			points[2] = RightOutputPoint
# 			points[3] = self.points[i]+TaperLength*array([cos(self.angles[i]+pi),sin(self.angles[i]+pi)])+vec*WaveguideWidth/2
# 		#	vecR = array([cos(self.angles[i-1]-pi/2),sin(self.angles[i-1]-pi/2)])
# 		#	P = points[2]-points[4]
# 		#	a = (P[1]*vec[0]-P[0]*vec[1])/(vec[0]*vecR[1]-vec[1]*vecR[0]) 
# 		#	points[3] = points[4]-vec*a	
# 			RightCollectionPoint = self.points[i]-0.5*self.widths[i]*vec
# 			RightOutputPoint = points[3]-vec*WaveguideWidth
# 		#	cell.add(gdspy.Round(2, self.points[i], 1e-6))
# 			cell.add(gdspy.Polygon(2, points))
# 		angles = range(-60,60,self.n)
# 		
# 	
# 		gdspy.gds_view()
	
